package cn.six1943.deepchat.module.ai.service.chatsession;

import cn.hutool.core.util.StrUtil;
import cn.six1943.deepchat.module.ai.agent.*;
import cn.six1943.deepchat.module.ai.controller.admin.chatsession.vo.ChatRequest;
import com.alibaba.fastjson.JSON;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.service.AiServices;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

/**
 *
 */
@AllArgsConstructor
@Service(AiConstant.agent.SERVICE_PLAINMODEL)
@Slf4j
public class AgentPlainModelServiceImpl implements AgentService<ChatRequest> {

    // private final AiProperties props;
    // private final AssistantMapper assistantMapper;

    @Override
    public void answerTheQuestion(ChatRequest request, ChatHandler handler) {
        log.info("开始问答，{}", JSON.toJSONString(request));
        StreamingChatModel chatModel = ModelHolder.getChatModel(request.getAgentId(), request.getModel());
        Assistant assistant = AiServices.builder(Assistant.class)
                .streamingChatModel(chatModel)
                .build();
        List<AgentOutput.ModelUsage> modelUsages = new ArrayList<>();
        long start = System.currentTimeMillis();
        log.info(StrUtil.format("agent::大模型问答开始::{}::{}", request.getSessionId(), start));
        assistant.answer(request.getContent())
                .onRetrieved(contents -> {
                })
                .onPartialResponse(token -> {
                    log.info(StrUtil.format("agent::大模型问答进行::{}::耗时{}", request.getSessionId(), System.currentTimeMillis() - start));
                    handler.onChunk(request.getSessionId(), request.getMessageId(), token);
                })
                .onCompleteResponse(resp -> {
                    log.info(StrUtil.format("agent::大模型问答结束::{}::{}::耗时{}", request.getSessionId(), System.currentTimeMillis(), System.currentTimeMillis() - start));
                    TokenUsage usage = resp.tokenUsage();
                    AgentOutput.ModelUsage.ModelUsageBuilder modelUsage = AgentOutput.ModelUsage.builder()
                            .modelId(String.valueOf(request.getAgentId()))
                            .inputTokens(usage.inputTokenCount())
                            .outputTokens(usage.outputTokenCount())
                            .totalTokens(usage.totalTokenCount());
                    modelUsages.add(modelUsage.build());
                    AgentOutput output = AgentOutput.builder()
                            .models(modelUsages)
                            .build();
                    handler.onComplete(request.getSessionId(), request.getMessageId(), output);
                })
                .onError(e -> {
                    log.info(StrUtil.format("大模型问答报错::{}::{}", request.getSessionId(), System.currentTimeMillis()));
                    log.error("model api error", e);
                    handler.onError(request.getSessionId(), request.getMessageId(), e.getMessage());
                })
                .start();
    }
}
